// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#define XXH_INLINE_ALL
#include <xxhash.h> // @manual=third-party//xxHash:xxhash

#include <pollux/exec/aggregate.h>
#include <pollux/expression/function_signature.h>
#include <pollux/functions/prestosql/aggregates/aggregate_names.h>
#include <pollux/functions/prestosql/aggregates/checksum_aggregate.h>
#include <pollux/functions/prestosql/aggregates/presto_hasher.h>
#include <pollux/vector/flat_vector.h>

namespace kumo::pollux::aggregate::prestosql {
    namespace {
        /// Computes an order-insensitive checksum of the input
        /// vector.
        ///
        /// checksum(T)-> varbinary
        class ChecksumAggregate : public exec::Aggregate {
        public:
            explicit ChecksumAggregate(const TypePtr &resultType)
                : Aggregate(resultType) {
            }

            int32_t accumulatorFixedWidthSize() const override {
                return sizeof(int64_t);
            }

            void extractValues(char **groups, int32_t numGroups, VectorPtr *result)
            override {
                auto *vector = (*result)->as_unchecked<FlatVector<StringView> >();
                POLLUX_CHECK(vector);
                vector->resize(numGroups);

                auto *rawValues = vector->mutableRawValues();
                vector->clear_all_nulls();
                for (auto i = 0; i < numGroups; ++i) {
                    auto group = groups[i];
                    if (isNull(group)) {
                        vector->set_null(i, true);
                    } else {
                        rawValues[i] = StringView(value<char>(group), sizeof(int64_t));
                    }
                }
            }

            void extractAccumulators(char **groups, int32_t numGroups, VectorPtr *result)
            override {
                auto vector = (*result)->as_unchecked<FlatVector<int64_t> >();
                POLLUX_CHECK(vector);
                vector->resize(numGroups);

                auto *rawValues = vector->mutableRawValues();
                vector->clear_all_nulls();
                for (auto i = 0; i < numGroups; ++i) {
                    auto group = groups[i];
                    if (isNull(group)) {
                        vector->set_null(i, true);
                    } else {
                        rawValues[i] = *value<int64_t>(group);
                    }
                }
            }

            bool isNullOrNullArray(const TypePtr &type) {
                if (type->is_unKnown()) {
                    return true;
                }
                // Only supports null array type for now.
                if (type->kind() == TypeKind::ARRAY) {
                    return type->as_array().elementType()->is_unKnown();
                }
                return false;
            }

            void addRawInput(
                char **groups,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool /*mayPushDown*/) override {
                const auto &arg = args[0];

                if (isNullOrNullArray(arg->type())) {
                    rows.applyToSelected([&](auto row) {
                        auto group = groups[row];
                        clearNull(group);
                        if (arg->is_null_at(row)) {
                            computeHashForNull(group);
                        } else {
                            computeHashForEmptyNullArray(group);
                        }
                    });
                    return;
                }

                auto hasher = getPrestoHasher(arg->type());
                auto hashes = getHashBuffer(rows.end(), arg->pool());
                hasher->hash(arg, rows, hashes);
                auto rawHashes = hashes->as<int64_t>();

                rows.applyToSelected([&](vector_size_t row) {
                    auto group = groups[row];
                    clearNull(group);
                    if (arg->is_null_at(row)) {
                        computeHashForNull(group);
                    } else {
                        computeHash(group, rawHashes[row]);
                    }
                });
            }

#if defined(MELON_DISABLE_UNDEFINED_BEHAVIOR_SANITIZER)
            MELON_DISABLE_UNDEFINED_BEHAVIOR_SANITIZER("signed-integer-overflow")
#endif
            MELON_ALWAYS_INLINE void safeAdd(int64_t &lhs, const int64_t &rhs) {
                lhs += rhs;
            }

            void addIntermediateResults(
                char **groups,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool /*mayPushDown*/) override {
                decodedIntermediate_.decode(*args[0], rows);

                rows.applyToSelected([&](vector_size_t row) {
                    auto group = groups[row];
                    if (!decodedIntermediate_.is_null_at(row)) {
                        clearNull(group);
                        safeAdd(
                            *value<int64_t>(group), decodedIntermediate_.value_at<int64_t>(row));
                    }
                });
            }

            void addSingleGroupRawInput(
                char *group,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool /*mayPushDown*/) override {
                const auto &arg = args[0];

                if (isNullOrNullArray(arg->type())) {
                    rows.applyToSelected([&](auto row) {
                        clearNull(group);
                        if (arg->is_null_at(row)) {
                            computeHashForNull(group);
                        } else {
                            computeHashForEmptyNullArray(group);
                        }
                    });

                    return;
                }

                auto hasher = getPrestoHasher(arg->type());
                auto hashes = getHashBuffer(rows.end(), arg->pool());
                hasher->hash(arg, rows, hashes);
                auto rawHashes = hashes->as<int64_t>();

                rows.applyToSelected([&](vector_size_t row) {
                    clearNull(group);
                    if (arg->is_null_at(row)) {
                        computeHashForNull(group);
                    } else {
                        computeHash(group, rawHashes[row]);
                    }
                });
            }

            void addSingleGroupIntermediateResults(
                char *group,
                const SelectivityVector &rows,
                const std::vector<VectorPtr> &args,
                bool /*mayPushDown*/) override {
                decodedIntermediate_.decode(*args[0], rows);

                int64_t result = 0;
                bool clearGroupNull = false;
                rows.applyToSelected([&](vector_size_t row) {
                    if (!decodedIntermediate_.is_null_at(row)) {
                        safeAdd(result, decodedIntermediate_.value_at<int64_t>(row));
                        clearGroupNull = true;
                    }
                });
                if (clearGroupNull) {
                    clearNull(group);
                }
                safeAdd(*value<int64_t>(group), result);
            }

        protected:
            void initializeNewGroupsInternal(
                char **groups,
                melon::Range<const vector_size_t *> indices) override {
                setAllNulls(groups, indices);
                for (auto i: indices) {
                    *value<int64_t>(groups[i]) = 0;
                }
            }

        private:
            MELON_ALWAYS_INLINE void computeHash(char *group, const int64_t hash) {
                *value<int64_t>(group) += hash * XXH_PRIME64_1;
            }

            MELON_ALWAYS_INLINE void computeHashForNull(char *group) {
                *value<int64_t>(group) += XXH_PRIME64_1;
            }

            MELON_ALWAYS_INLINE void computeHashForEmptyNullArray(char *group) {
                *value<int64_t>(group) += 0;
            }

            MELON_ALWAYS_INLINE PrestoHasher *getPrestoHasher(TypePtr typePtr) {
                if (prestoHasher_ == nullptr) {
                    prestoHasher_ = std::make_unique<PrestoHasher>(std::move(typePtr));
                }
                return prestoHasher_.get();
            }

            MELON_ALWAYS_INLINE BufferPtr &getHashBuffer(
                vector_size_t size,
                pollux::memory::MemoryPool *pool) {
                // hashes_->size() is in bytes.
                if (hashes_ == nullptr) {
                    hashes_ = AlignedBuffer::allocate<int64_t>(size, pool);
                } else if (hashes_->size() < (size * sizeof(int64_t))) {
                    AlignedBuffer::reallocate<int64_t>(&hashes_, size);
                }

                return hashes_;
            }

            std::unique_ptr<PrestoHasher> prestoHasher_;
            BufferPtr hashes_;
            DecodedVector decodedIntermediate_;
        };
    } // namespace

    void registerChecksumAggregate(
        const std::string &prefix,
        bool withCompanionFunctions,
        bool overwrite) {
        std::vector<std::shared_ptr<exec::AggregateFunctionSignature> > signatures{
            exec::AggregateFunctionSignatureBuilder()
            .typeVariable("T")
            .returnType("varbinary")
            .intermediateType("bigint")
            .argumentType("T")
            .build(),
        };

        auto name = prefix + kChecksum;
        exec::registerAggregateFunction(
            name,
            std::move(signatures),
            [&name](
        core::AggregationNode::Step step,
        const std::vector<TypePtr> &argTypes,
        const TypePtr & /*resultType*/,
        const core::QueryConfig & /*config*/)
        -> std::unique_ptr<exec::Aggregate> {
                POLLUX_CHECK_EQ(argTypes.size(), 1, "{} takes one argument", name);

                if (exec::isPartialOutput(step)) {
                    return std::make_unique<ChecksumAggregate>(BIGINT());
                }

                return std::make_unique<ChecksumAggregate>(VARBINARY());
            },
            {false /*orderSensitive*/, false /*companionFunction*/},
            withCompanionFunctions,
            overwrite);
    }
} // namespace kumo::pollux::aggregate::prestosql
